from typing import Any, Dict, List
import numpy as np
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.llms import DeepSeek
from sentence_transformers import CrossEncoder
from ..base import BaseFusionModule, SearchResult, Document


class QueryExpander:
    """查询扩展器"""
    def __init__(self, llm):
        self.llm = llm
        self.template = """基于以下原始查询，生成多个扩展查询，以获取更全面的信息：

原始查询: {query}

请生成不同角度的查询，考虑：
1. 同义词和相关术语
2. 更具体的子查询
3. 更广泛的上下文
4. 潜在的关联主题

生成的查询："""
        
        self.prompt = PromptTemplate(
            template=self.template,
            input_variables=["query"]
        )
        self.chain = LLMChain(llm=llm, prompt=self.prompt)

    async def expand(self, query: str) -> List[str]:
        """扩展查询"""
        result = await self.chain.arun(query=query)
        return [q.strip() for q in result.split("\n") if q.strip()]


class ResultMerger:
    """结果合并器"""
    def __init__(self, strategies: List[Dict[str, float]]):
        self.strategies = strategies

    def reciprocal_rank_fusion(self, results: List[SearchResult], k: float = 60.0) -> SearchResult:
        """倒数排名融合"""
        # 为每个文档计算RRF分数
        doc_scores = {}
        for rank, result in enumerate(results):
            for doc in result.documents:
                if doc.content not in doc_scores:
                    doc_scores[doc.content] = 0
                doc_scores[doc.content] += 1.0 / (rank + k)

        # 创建融合后的文档列表
        merged_docs = [
            Document(
                content=content,
                metadata={"score": score},
                score=score
            )
            for content, score in sorted(
                doc_scores.items(),
                key=lambda x: x[1],
                reverse=True
            )
        ]

        return SearchResult(
            documents=merged_docs,
            source="fusion",
            metadata={"strategy": "reciprocal_rank_fusion"}
        )

    def weighted_score_fusion(self, results: List[SearchResult]) -> SearchResult:
        """加权分数融合"""
        # 为每个文档计算加权分数
        doc_scores = {}
        for result, strategy in zip(results, self.strategies):
            weight = strategy["weight"]
            for doc in result.documents:
                if doc.content not in doc_scores:
                    doc_scores[doc.content] = 0
                doc_scores[doc.content] += (doc.score or 1.0) * weight

        # 创建融合后的文档列表
        merged_docs = [
            Document(
                content=content,
                metadata={"score": score},
                score=score
            )
            for content, score in sorted(
                doc_scores.items(),
                key=lambda x: x[1],
                reverse=True
            )
        ]

        return SearchResult(
            documents=merged_docs,
            source="fusion",
            metadata={"strategy": "weighted_score_fusion"}
        )


class Reranker:
    """重排序器"""
    def __init__(self):
        self.model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

    def rerank(self, query: str, results: SearchResult) -> SearchResult:
        """重排序结果"""
        # 准备文档对
        pairs = [[query, doc.content] for doc in results.documents]
        
        # 计算相关性分数
        scores = self.model.predict(pairs)
        
        # 结合文档和分数
        scored_docs = list(zip(results.documents, scores))
        
        # 按分数排序
        ranked_docs = [
            Document(
                content=doc.content,
                metadata=doc.metadata,
                score=score
            )
            for doc, score in sorted(
                scored_docs,
                key=lambda x: x[1],
                reverse=True
            )
        ]

        return SearchResult(
            documents=ranked_docs,
            source="rerank",
            metadata={"strategy": "cross_encoder"}
        )


class FusionModule(BaseFusionModule):
    """融合模块实现"""
    def __init__(self):
        self.llm = None
        self.query_expander = None
        self.result_merger = None
        self.reranker = None
        self.config = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化模块"""
        self.config = config
        self.llm = DeepSeek(
            model=config["model"]["name"],
            temperature=config["model"]["temperature"],
            max_tokens=config["model"]["max_tokens"]
        )
        
        self.query_expander = QueryExpander(self.llm)
        self.result_merger = ResultMerger(config["strategies"])
        self.reranker = Reranker()

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["model", "strategies", "max_results", "score_threshold"]
        if not all(key in config for key in required):
            return False
            
        strategy_required = ["name", "weight"]
        return all(
            all(key in strategy for key in strategy_required)
            for strategy in config["strategies"]
        )

    async def expand_query(self, query: str) -> List[str]:
        """扩展查询"""
        if not self.query_expander:
            raise RuntimeError("Module not initialized")
            
        expanded = await self.query_expander.expand(query)
        return expanded[:self.config["max_results"]]

    async def merge_results(self, results: List[SearchResult]) -> SearchResult:
        """合并结果"""
        if not self.result_merger:
            raise RuntimeError("Module not initialized")
            
        # 根据策略选择融合方法
        strategy_names = [s["name"] for s in self.config["strategies"]]
        if "reciprocal_rank" in strategy_names:
            merged = self.result_merger.reciprocal_rank_fusion(results)
        else:
            merged = self.result_merger.weighted_score_fusion(results)
            
        # 过滤低分结果
        filtered_docs = [
            doc for doc in merged.documents
            if doc.score >= self.config["score_threshold"]
        ]
        merged.documents = filtered_docs[:self.config["max_results"]]
        
        return merged

    async def rerank(self, query: str, results: SearchResult) -> SearchResult:
        """重排序结果"""
        if not self.reranker:
            raise RuntimeError("Module not initialized")
            
        reranked = self.reranker.rerank(query, results)
        
        # 限制结果数量
        reranked.documents = reranked.documents[:self.config["max_results"]]
        return reranked 